
#include "rtmp_parser.h"
#include "rtmp_message.h"
#include <cstdint>
#include <cstring>
#include <endian.h>
#include <memory>
#include <vector>


RtmpParser::RtmpParser()
{
    m_header_length_map[0] = 11;
    m_header_length_map[1] = 7;
    m_header_length_map[2] = 3;
    m_header_length_map[3] = 0;
	// default chunk size is 128
    m_chunk_size = 128;
}

RtmpParser::~RtmpParser()
{
}

std::shared_ptr<RtmpMessage> RtmpParser::ParseRtmpMessage(const uint8_t * data, int length, int & use_length, int status)
{
    (void)status;
    std::shared_ptr<RtmpMessage> rtmp_message;
    if (m_rtmp_msg)
    {
        rtmp_message = m_rtmp_msg;
    }
    else
    {
        rtmp_message = std::make_shared<RtmpMessage>();
		rtmp_message->m_header = {};
    }
    int index = 0;
    while(true)
    {
        int result = ProcessHeader(data + index, length - index, rtmp_message);
        if (result < 0)
        {
            rtmp_message = nullptr;
            m_rtmp_msg   = nullptr;
            m_data.clear();
            // drop all data
            use_length = length;
			WARN("parser error drop length : {}", use_length);
            break;
        }
        else
        {
            index += result;
            int header_length = result;
            result = ProcessBody(data + index, length - index, rtmp_message);
            if (result == 0)
            {
                m_rtmp_msg = rtmp_message;
                rtmp_message = nullptr;
                use_length = index - header_length;
                break;
            }
            else
            {
                index += result;
                if (rtmp_message->m_header.m_body_size == m_data.size())
                { 
                    result = SwitchBodyCommand(rtmp_message);
                    if (result == -1)
                    {
                        m_rtmp_msg = nullptr;
                        rtmp_message = nullptr;
                        m_data.clear();
                        use_length = length;
                    }
                    else
                    {
                        // parse finish clear tmp data
                        m_rtmp_msg = nullptr;
                        m_data.clear();
                        use_length = index; 
                    }
                    break;
                }
            }
        }
    }
    return rtmp_message;
}

int RtmpParser::SwitchBodyCommand(std::shared_ptr<RtmpMessage> & message)
{
    int result = 0; 
    switch(message->m_header.m_type_id)
    {
        case 0x01:
        {
            uint32_t raw_data = BeToH32(&*(m_data.begin()),m_data.size());
            m_chunk_size = raw_data;
            INFO("new chunk size : {}", m_chunk_size);
            result = message->m_header.m_body_size;
            message->m_body.m_attr_data.m_chunk_size = m_chunk_size;
            break;
        }
        case 0x05:
        {
            /* result = message->m_header.m_body_size; */
            result = 1;
            break; 
        }
        case 0x06:
        {
            result = message->m_header.m_body_size;
            break;
        }
        case 0x14:
        {
            // command message 
            result = ParseBodyData(message);
            break;
        }
        case 0x12:
        {
            // metadata frame
            result = ParseMetaData(message);
            break;
        }
        case 0x08:
        {
            // audio frame data
            result = ParseAudioData(message);
            break;
        }
        case 0x09:
        {
            // video frame data
            result = ParseVideoData(message);
            break;
        }
        case 0x04:
        {
            if (m_data.size() >= message->m_header.m_body_size)
            {
                result = message->m_header.m_body_size;
            }
            else
            {
                result = 0;
            }
            break;
        }
        default:
        {
            result = -1;
        }
    }
    return result;
}

int RtmpParser::ProcessHeader(const uint8_t * data, int length, std::shared_ptr<RtmpMessage> & message)
{
    int index = 0; 
    uint8_t basic_head = data[index++];
    uint8_t fmt = (basic_head >> 6) & 0x03;
    if (message->m_header.m_format == 0)
    {
        message->m_header.m_format = fmt;
    }
    if (message->m_header.m_chunk_stream == 0)
    {
        message->m_header.m_chunk_stream = basic_head & 0x3f;
    }
    if (m_header_length_map.find(fmt) == m_header_length_map.end())
    {
        return -1;
    }
    int header_length = m_header_length_map[fmt];
    if (header_length > length)
    {
        return 0;
    }
    if (header_length >= 3)
    {
        uint32_t timestamp = 0; 
        memcpy(&timestamp, data + index, 3);
        message->m_header.m_timestamp = BeToHn((uint8_t*)&timestamp, sizeof(timestamp), 3);
        index += 3;
    }
    if (header_length >= 7)
    { 
        uint32_t body_size = 0;
        memcpy(&body_size, data + index, 3);
        message->m_header.m_body_size = BeToHn((uint8_t*)&body_size, sizeof(body_size), 3);

        index += 3;
        message->m_header.m_type_id = *(data + index);
        index++;
    }
    if (header_length >= 11)
    {
        uint32_t stream_id = 0;
        memcpy(&stream_id, data + index, 4);
        message->m_header.m_stream_id = stream_id;
        index += 4;
    }
	if (message->m_header.m_body_size == 0)
	{
		auto iter = m_rtmp_stream_map.find(message->m_header.m_chunk_stream);
		if (iter != m_rtmp_stream_map.end())
		{
			message->m_header.m_body_size = iter->second.m_record_body_size;
			message->m_header.m_type_id   = iter->second.m_type_id;
			if (header_length == 0)
			{
				message->m_header.m_timestamp = iter->second.m_delta_time;
			}
			// TRACE("body size : {} timestamp : {} type id : {} header length : {}", message->m_header.m_body_size, message->m_header.m_timestamp, message->m_header.m_type_id, header_length); 
		}
		else
		{
			WARN("no found streamchunk id : {}", message->m_header.m_chunk_stream);
		}
	}
	return index;
}

int RtmpParser::ProcessBody(const uint8_t * data, int length, std::shared_ptr<RtmpMessage> & message)
{
    int index = 0;
    int need_size = message->m_header.m_body_size - m_data.size();
    int actual_size = need_size >= (int)m_chunk_size ? m_chunk_size : need_size; 
    if (length >= actual_size)
    {
        for (int i = 0; i < actual_size; i++)
        {
            m_data.push_back(*(data + i));
        }
        index += actual_size;
    }
    return index;
}

int RtmpParser::ParseBodyData(std::shared_ptr<RtmpMessage> & message)
{
    int index = 0;
    for (;index < (int)m_data.size();)
    { 
        int use_size = 0;
        auto value = AmfValue::ParseBuffer(&*(m_data.begin() + index), m_data.size() - index, use_size);
        if (value)
        { 
            message->m_body.m_amf_value_vec.push_back(value);
            index += use_size;
        }
        else
        {
            index = -1;
            break;
        }
    }
    return index;
}

int RtmpParser::ParseMetaData(std::shared_ptr<RtmpMessage> & message)
{

    return ParseBodyData(message);
}

int RtmpParser::ParseAudioData(std::shared_ptr<RtmpMessage> & message)
{
	auto iter = m_rtmp_stream_map.find(message->m_header.m_chunk_stream);
	if (iter == m_rtmp_stream_map.end())
	{
		RtmpStreamData data{};
		data.m_stream_id        = message->m_header.m_chunk_stream;
		data.m_record_body_size = message->m_header.m_body_size;
		data.m_timestamp        = message->m_header.m_timestamp;
		data.m_type_id          = message->m_header.m_type_id;
		data.m_delta_time       = message->m_header.m_timestamp;
		message->m_header.m_timestamp       = data.m_timestamp;
		message->m_header.m_calc_timestamp  = data.m_timestamp;
		m_rtmp_stream_map[data.m_stream_id] = data;
		TRACE("add new stream id {} and body size : {} delta timestamp : {} type id : {} timestamp : {}", data.m_stream_id, data.m_record_body_size, data.m_delta_time, data.m_type_id, data.m_timestamp);
   	}
	else
	{
		iter->second.m_timestamp += message->m_header.m_timestamp;
		iter->second.m_record_body_size = message->m_header.m_body_size;
		message->m_header.m_calc_timestamp = iter->second.m_timestamp;
   	}
    int index = 0;
    message->m_body.m_attr_data.m_media_attr_data = m_data[index++];
    for (int i = 1; i < (int)m_data.size(); i++)
    {
        message->m_body.m_data.push_back(m_data[index++]);
    }
    return index;
}

int RtmpParser::ParseVideoData(std::shared_ptr<RtmpMessage> & message)
{
    int index = ParseAudioData(message);
    return index;
}

int RtmpParser::GetChunkSize()
{
    return (int)m_chunk_size;
}

RtmpMessageCompose::RtmpMessageCompose(RtmpParser * parser)
{
    m_rtmp_parser = parser;
    m_header_length_map[0] = 11;
    m_header_length_map[1] = 7;
    m_header_length_map[2] = 3;
    m_header_length_map[3] = 0;
    /* m_chunk_size = 128; */
    m_chunk_size = m_rtmp_parser->GetChunkSize();
}

RtmpMessageCompose::~RtmpMessageCompose()
{
}

int RtmpMessageCompose::ComposeRtmpMessage(std::shared_ptr<RtmpMessage> & msg)
{
    int length = 0;
    for (auto & data : msg->m_body.m_amf_value_vec)
    {
        length += data->GetLength();
    }
    if (!msg->m_body.m_data.empty())
    { 
        length++;
        length += msg->m_body.m_data.size();
    }
    else if (msg->m_body.m_amf_value_vec.empty())
    {
        // other message command use body size
        length = msg->m_header.m_body_size; 
    }
    int index = 0;
    std::vector<uint8_t> body_buffer(length);
    for (auto & data : msg->m_body.m_amf_value_vec)
    {
        index += data->CopyToBuffer(&*(body_buffer.begin() + index), length - index);
    }
    if (!msg->m_body.m_data.empty())
    {
        body_buffer.push_back(msg->m_body.m_attr_data.m_media_attr_data);
        memcpy(&*(body_buffer.begin() + 1),&*msg->m_body.m_data.begin(), msg->m_body.m_data.size());
    }
    else if (msg->m_body.m_amf_value_vec.empty())
    {
        ComposeRtmpMessageAttrData(msg, body_buffer);
    }
    msg->m_header.m_body_size = length;
    int chunk_count = length/m_chunk_size + (length % m_chunk_size == 0 ? 0 : 1);
    auto copy_header = msg->m_header;
    for (int i = 0; i < chunk_count; i++)
    { 
        ComposeHeader(copy_header, i);
        int position = i * m_chunk_size;
        int chunk_body_length = length - position > m_chunk_size ? m_chunk_size : (length - position);
        WriteBuffer & write_buffer = m_index_buffer[i];
        for (int j = 0; j < chunk_body_length; j++)
        {
            write_buffer.WriteUint8(body_buffer.at(position + j));
        }
    }
    return 0;
}

void RtmpMessageCompose::ComposeRtmpMessageAttrData(std::shared_ptr<RtmpMessage> & msg, std::vector<uint8_t> & buffer)
{
    switch(msg->m_header.m_type_id)
    {
        case 0x01:
        {
            uint32_t data = msg->m_body.m_attr_data.m_chunk_size;
            memcpy(&*(buffer.begin()), &data, sizeof(data));
            break;
        } 
        case 0x02:
        {
            break;
        }
    }
}

int RtmpMessageCompose::ChunkBlockSize()
{

    return (int)m_index_buffer.size();
}

void RtmpMessageCompose::GetData(std::vector<uint8_t> & buffer)
{
    for (auto & data : m_index_buffer)
    {
        int length = data.second.GetBufferLength();
        uint8_t * ptr = data.second.GetBufferData();
        for (int index = 0; index < length; index++)
        {
            buffer.push_back(*(ptr + index));
        }
    }
    m_index_buffer.clear();
}

int RtmpMessageCompose::ComposeHeader(RtmpHeader & header, int index)
{
    m_index_buffer[index];
    WriteBuffer & write_buffer = m_index_buffer[index];
    uint8_t fmt = header.m_format;
    auto iter = m_header_length_map.find(fmt);
    int header_lnegth = iter->second;
    
    write_buffer.WriteUint8((fmt << 6) | header.m_chunk_stream);
    if (header_lnegth >= 3)
    {
        if (header.m_timestamp != 0)
        {
            write_buffer.WriteUint8(header.m_timestamp >> 16 & 0xff);
            write_buffer.WriteUint8(header.m_timestamp >> 8  & 0xff);
            write_buffer.WriteUint8(header.m_timestamp & 0xff);

            header.m_timestamp = 0;
        }
		else
		{
			write_buffer.WriteUint8(00 & 0xff);
			write_buffer.WriteUint8(00 & 0xff);
			write_buffer.WriteUint8(00 & 0xff);
		}

    }
    if (header_lnegth >= 7)
    {
        write_buffer.WriteUint8(header.m_body_size >> 16 & 0xff);
        write_buffer.WriteUint8(header.m_body_size >> 8  & 0xff);
        write_buffer.WriteUint8(header.m_body_size & 0xff);
        write_buffer.WriteUint8(header.m_type_id);

        header.m_body_size = 0;
        header.m_type_id   = 0;
    }
    if (header_lnegth >= 11)
    {
       write_buffer.WriteUint32(htobe32(header.m_stream_id));
    }
    switch(fmt)
    {
        case 0:
        {
            header.m_format = 3;
			break;
        }
        case 1:
        {
            header.m_format = 3;
			break;
        }
        case 2:
        {
            header.m_format = 3;
			break;
        }
        case 3:
        {
            break;
        }
    }
    return index;
}


RtmpMessageBuffer::RtmpMessageBuffer()
{
    m_compose = new RtmpMessageCompose(&m_parser);
}

RtmpMessageBuffer::~RtmpMessageBuffer()
{
    if (m_compose)
    {
        delete m_compose;
        m_compose = nullptr;
    }
}

int RtmpMessageBuffer::PushDataToBuffer(const uint8_t * data, int length, int & use_length)
{
    int index = 0;
    do 
    {
		int status   = 0;
		int use_size = 0;
		auto msg = m_parser.ParseRtmpMessage(data + use_length, length - use_length, use_size, status);
		(void)status;
		if (msg)
		{
			m_message_vec.push_back(msg);
			use_length += use_size;
			index++;
		}
		else
		{
			if ((length - use_length) >= use_size)
			{
				use_length += use_size;
				break;
			}
			else
			{
				use_length += use_size;
			}
		}
		if (use_length == length)
		{
			break;
		}
    }
    while(true);
    return index;
}

int RtmpMessageBuffer::MsgSize()
{
    return (int)m_message_vec.size();
}

std::shared_ptr<RtmpMessage> RtmpMessageBuffer::PopMsg()
{
    if (m_message_vec.size() == 0)
    {
        return nullptr;
    }
    auto msg = m_message_vec.front();
    m_message_vec.erase(m_message_vec.begin());
    return msg;
}

int RtmpMessageBuffer::ComposeRtmpMessage(std::shared_ptr<RtmpMessage> & msg)
{
    return m_compose->ComposeRtmpMessage(msg);
}

void RtmpMessageBuffer::GetData(std::vector<uint8_t> & buffer)
{
    m_compose->GetData(buffer);
}


void RtmpPacketDump::RtmpPacketDumpInfo(std::shared_ptr<RtmpMessage> packet)
{

}    
    
